#!/usr/bin/env python
from __future__ import print_function
import os
import sys
import time
import shlex
import datetime
import subprocess
from multiprocessing import Pool

os.environ['OMP_NUM_THREADS'] = '1' # do not use multithreading

CWD_DEFAULT = os.getcwd()

def run_script(shell, cwd, script, idx):
    print(' [%i]  Running %s on pid %s ...' % (idx, repr(script), os.getpid()), file=sys.stderr)

    if script.endswith('&'):
        script = script[:-1].strip()

    # pipe flags
    narr= script.count('&>>')
    nar = script.count('&>') - narr
    na  = script.count('&') - nar - narr
    n2rr= script.count('2>>')
    nrr = script.count('>>') - narr - n2rr
    n2r = script.count('2>') - n2rr
    nr  = script.count('>') - 2*(nrr+narr+n2rr) - nar - n2r
    #print(na, narr, nar, n2rr, nrr, n2r, nr)

    # extract infos
    if shell or na>0 or narr>1 or n2rr>1 or nar>1 or nrr>1 or n2r>1 or nr>1 or \
            (narr>0 and nar>0) or (nrr>0 and nr>0) or (n2r>0 and n2rr>0) or \
            ((nar>0 or narr>0) and (nrr>0 or n2r>0 or nr>0 or n2rr>0)):
        shell_flag = True
        mysc = script
    else:
        shell_flag = False
        logfile, errfile = None, None
        if narr == 1:
            mode1, mode2 = 'a', 'a'
            mysc, logfile = script.split('&>>')
            errfile = logfile
        elif nar == 1:
            mode1, mode2 = 'w', 'w'
            mysc, logfile = script.split('&>')
            errfile = logfile
        elif n2rr == 1:
            mode2 = 'a'
            p1, p2 = script.split('2>>')
            if nrr == 1:
                mode1 = 'a'
                if '>>' in p1:
                    scspl1 = p1.split('>>')
                    mysc = scspl1[0]
                    logfile = scspl1[1]
                    errfile = p2
                else:
                    mysc = p1
                    scspl1 = p2.split('>>')
                    errfile = scspl1[0]
                    logfile = scspl1[1]
            elif nr == 1:
                mode1 = 'w'
                if '>' in p1:
                    scspl1 = p1.split('>')
                    mysc = scspl1[0]
                    logfile = scspl1[1]
                    errfile = p2
                else:
                    mysc = p1
                    scspl1 = p2.split('>')
                    errfile = scspl1[0]
                    logfile = scspl1[1]
            else:
                mysc = p1
                errfile = p2
        elif nrr == 1:
            mode1 = 'a'
            p1, p2 = script.split('>>')
            if n2r == 1:
                mode2 = 'w'
                if '2>' in p1:
                    scspl1 = p1.split('2>')
                    mysc = scspl1[0]
                    errfile = scspl1[1]
                    logfile = p2
                else:
                    mysc = p1
                    scspl1 = p2.split('2>')
                    logfile = scspl1[0]
                    errfile = scspl1[1]
            else:
                mysc = p1
                logfile = p2
        elif n2r == 1:
            mode2 = 'w'
            p1, p2 = script.split('2>')
            if nr == 1:
                mode1 = 'w'
                if '>' in p1:
                    scspl1 = p1.split('>')
                    mysc = scspl1[0]
                    logfile = scspl1[1]
                    errfile = p2
                elif '>' in p2:
                    mysc = p1
                    scspl1 = p2.split('>')
                    errfile = scspl1[0]
                    logfile = scspl1[1]
            else:
                mysc = p1
                errfile = p2
        elif nr == 1:
            mode1 = 'w'
            mysc, logfile = script.split('>')
        else:
            mysc = script

        #print(mysc)
        #print(logfile, errfile)
        mysc = mysc.strip()
        if logfile is not None:
            logfile = logfile.strip()
        if errfile is not None:
            errfile = errfile.strip()

    if shell_flag:
        ret = subprocess.call(mysc, cwd=cwd, shell=True)
    else:
        myscspl = shlex.split(mysc)

        # prepare logfile handle
        if logfile is None:
            logfh = None
            logfh_open = False
        else:
            if logfile.startswith('/'):
                mylogfile = logfile
            else:
                mylogfile = '%s/%s'%(cwd, logfile)
            try:
                logfh = open(mylogfile, mode1, 1)
            except:
                raise RuntimeError('Command %r fail to open %r' % (script, mylogfile))
            logfh_open = True

        # prepare errfile handle
        if errfile is None:
            errfh = None
            errfh_open = False
        elif errfile == logfile:
            errfh = subprocess.STDOUT
            errfh_open = False
        else:
            if errfile.startswith('/'):
                myerrfile = errfile
            else:
                myerrfile = '%s/%s'%(cwd, errfile)
            try:
                errfh = open(myerrfile, mode2, 1)
            except:
                raise RuntimeError('Command %r fail to open %r' % (script, myerrfile))
            errfh_open = True

        # run the script
        try:
            ret = subprocess.call(myscspl, cwd=cwd, stderr=errfh, stdout=logfh)
        except OSError:
            ret = 1
        finally:
            if logfh_open:
                logfh.close()
            if errfh_open:
                errfh.close()

    if ret != 0:
        raise RuntimeError('Command %r returned non-zero exit status %i' % (script, ret))


def mpi_from_joblist(joblists, nthread=5, shell=False, raiseErr=True, errorFile=None):
    """
    joblist is arrange like this:
        DIRECTORY1 : SCRIPT1 # COMMENT1
        DIRECTORY2 : SCRIPT2 # COMMENT2
        ...
        DIRECTORYn : SCRIPTn # COMMENTn
    For sake of security, the job containing 'rm' will be skipped.
    """
    pool_list = []
    for joblist in joblists:
        print(' loading the joblist %s ...' % joblist, file=sys.stderr)
        with open(joblist, 'r') as fh:
            for string in fh:
                job = string.split('#')[0].strip()

                if len(job) < 1:
                    continue
                if 'rm -' in job:
                    print(' find "rm -" in script %s, skip the script' % job, file=sys.stderr)
                    continue
                if not ':' in job:
                    cwd = CWD_DEFAULT
                    scripts = job.strip()
                else:
                    parser = job.split(':')
                    cwd, scripts = parser[0].strip(), parser[1].strip()
                pool_list.append((cwd, scripts))

    njobs = len(pool_list)
    nthread = min(nthread, njobs)
    print(" submitting the jobs to %i threads ..." % nthread, file=sys.stderr)
    if nthread < 2:
        tst = time.time()
        next_ms, prev_tms = 0.05, tst

        for iap, (cwd, scripts) in enumerate(pool_list):
            # show some info
            if iap>njobs*next_ms:
                this_tms = time.time()
                trun = this_tms - tst
                pct = iap*1./njobs
                while next_ms<=pct:
                    next_ms += 0.16
                if this_tms-prev_tms > 10.:
                    prev_tms = this_tms
                    tfinish = this_tms + trun*(1.-pct)/pct
                    tfinish_utc = datetime.datetime.fromtimestamp(tfinish).strftime('%Y-%m-%dT%H:%M:%S')
                    infostr = ' | {:.0%}, Trun = {:.1f}[s], ExpFinish @{!s} |'.format(pct, trun, tfinish_utc)
                    wrapper = ' +' + '-'*(len(infostr)-3) + '+'
                    print(wrapper+'\n'+infostr+'\n'+wrapper, file=sys.stderr)

            # run
            #apply(run_script, (shell, cwd, scripts, iap))
            run_script(shell, cwd, scripts, iap)
    else:
        myapplies = []
        error_lines = []
        pool = Pool(processes=nthread)
        pool_names = [p.name for p in pool._pool]
        try:
            for iap, (cwd, scripts) in enumerate(pool_list):
                myapply = pool.apply_async(run_script, (shell, cwd, scripts, iap))
                myapplies.append(myapply)
            pool.close()

            nfinish = 0
            ready_list = [0]*njobs
            idx_start = 0
            tst = time.time()
            next_ms, prev_tms = 0.05, tst

            while True:
                for pw in pool._pool:
                    if not pw.is_alive() and pw.exitcode != 0:
                        raise RuntimeError(' %s is dead!' % pw.name, file=sys.stderr)
                    elif not (pw.name in pool_names):
                        raise RuntimeError(' unknown %s!' % pw.name, file=sys.stderr)

                n_undone = 0
                for iap, ap in enumerate(myapplies[idx_start:], idx_start):
                    if ready_list[iap]:
                        continue
                    else:
                        if ap.ready():
                            try:
                                ap.get()
                            except Exception as e:
                                emesg = str(e)
                                error_lines.append(pool_list[iap])
                                try:
                                    retcode = int(emesg.split('status')[1])
                                except:
                                    if raiseErr:
                                        print(' [%i]- Exit -1  %s'%(iap, repr(pool_list[iap][1])))
                                        raise e
                                    elif emesg:
                                        print(' [%i]- Exit -1  %s: %s'%(iap, repr(pool_list[iap][1]), emesg))
                                else:
                                    if raiseErr:
                                        print(' [%i]- Exit %i  %s' % (iap, retcode, repr(pool_list[iap][1])),
                                            file=sys.stderr)
                                        raise e
                                    else:
                                        print(' [%i]- Exit %i  %s: %s' % (iap, retcode, repr(pool_list[iap][1]), emesg),
                                            file=sys.stderr)
                            else:
                                print(' [%i]+ Done    %s'%(iap, repr(pool_list[iap][1])), file=sys.stderr)
                            ready_list[iap] = 1
                            nfinish += 1
                        else:
                            n_undone += 1
                            if n_undone == 1:
                                idx_start = iap
                            elif n_undone >= nthread:
                                break
                            continue

                if nfinish == njobs:
                    break

                # show some info
                if nfinish>njobs*next_ms:
                    this_tms = time.time()
                    trun = this_tms - tst
                    pct = nfinish*1./njobs
                    while next_ms <= pct:
                        next_ms += 0.16
                    if this_tms-prev_tms > 10.:
                    #if True:
                        prev_tms = this_tms
                        pct2 = (nfinish+n_undone/2.)/njobs
                        tfinish = this_tms + trun*(1.-pct2)/pct2 + 1.
                        tfinish_utc = datetime.datetime.fromtimestamp(tfinish).strftime('%Y-%m-%dT%H:%M:%S')
                        infostr = ' | {:.0%}, Trun = {:.1f}[s], ExpFinish @{!s} |'.format(pct, trun, tfinish_utc)
                        wrapper = ' +' + '-'*(len(infostr)-3) + '+'
                        print(wrapper+'\n'+infostr+'\n'+wrapper, file=sys.stderr)

                myapplies[idx_start].wait(1.)
        except:
            pool.terminate()
            raise

    print(" sub-process(es) done", file=sys.stderr)

    if errorFile is not None and error_lines:
        with open(errorFile, 'w') as fh2:
            for cwd, cmd in error_lines:
                fh2.write('%s: %s\n'%(cwd, cmd))

        print(" Error lines are saved to {!r}".format(errorFile))

    print(" Current time: {!s}".format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')), file=sys.stderr)


def cli():
    import argparse

    parser = argparse.ArgumentParser(description='The script is used to run scripts in a multiprocess way.', prefix_chars='-+')
    parser.add_argument("jobs", type=int, help="The number of jobs you wish to spawn (usually the number of cores on your machine).")
    parser.add_argument("joblists", nargs='+', help="The ASCII files containing a list of jobs which aranged as doc says.")
    parser.add_argument("+shell", action='store_true', default=False, help='Some script may need shell, add this flag (default: False).')
    parser.add_argument("+noerror", action='store_true', default=False, help='Whether to ignore the error (default: False).')
    parser.add_argument("-errfile", type=str, default=None, help='The pool file that contains the line that error occurs (default: None)')

    args = parser.parse_args()

    mpi_from_joblist(args.joblists, args.jobs, args.shell, raiseErr=(not args.noerror), errorFile=args.errfile)


if __name__ == '__main__': cli()
